Separate the routing for control flows and data flows #1082

Yu-Zhewen · 2025-02-06T00:13:06Z

The routing results for control flows should remain unchanged, no matter whether (or how many data flows) are present. Therefore, the routing pass is now run twice, with the first time to route control flows only, and the second to route data flows. To achieve this,

router can now load existing, partially routed packet flows from IR (previously, only circuit flows could be loaded).
the above information is passed to the Dijkstra pathfinding algorithm and the arbiter/MSel allocation algorithm, for deterministic behavior.

Note: mlir-aie handled the problem by still running the routing pass only once, and sorting flows before invoking Dijkstra's algorithm (https://github.com/Xilinx/mlir-aie/blob/a60888210c76266cc932d5a0cc922fceef0322f0/lib/Dialect/AIE/Transforms/AIEPathFinder.cpp#L349). However, since Dijkstra is an iterative, congestion-aware algorithm, sorting alone does not fully prevent congestion from data flows affecting control flow routing.

Yu-Zhewen · 2025-02-11T23:18:03Z

compiler/plugins/target/AMD-AIE/aie/AIEOps.td

@@ -236,7 +236,8 @@ def AIE_PacketRuleOp: AIE_Op<"rule"> {
  let arguments = (
    ins I8Attr:$mask,
        I8Attr:$value,
-        Index:$amsel
+        Index:$amsel,
+        OptionalAttr<DenseI32ArrayAttr>:$packet_ids


When lowering packet_flow to packet_rule, the IDs can be used to calculate the mask and value. However, the reverse operation (calculating the IDs from the mask and value) is not feasible, as we will probably get a superset of actual IDs. This can cause issues during the second run of the router.

Yu-Zhewen · 2025-02-11T23:20:32Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

@@ -108,22 +139,17 @@ PacketRulesOp getOrCreatePacketRules(OpBuilder &builder, SwitchboxOp &swboxOp,
  return packetRules;
 }

-struct ConvertFlowsToInterconnect : OpConversionPattern<FlowOp> {


Can't use OpConversionPattern, as it is also conditional based on whether it is control flow or not.

Yu-Zhewen · 2025-02-11T23:34:31Z

compiler/plugins/target/AMD-AIE/aie/test/test_pktflow_weight_pusher.mlir

-// CHECK:               aie.rule(31, 10, %3)
-// CHECK:               aie.rule(27, 11, %3)
-// CHECK:               aie.rule(24, 8, %2)
+// CHECK:               aie.rule(31, 13, %[[VAL_70]])


Due to the changes on
https://github.com/nod-ai/iree-amd-aie/pull/1082/files#diff-b5e43c7f6b1588e5d27b15fcfcc885d11c8d9753b8b1c93448ecf1ba3ee11bbcL849-R947

Yu-Zhewen · 2025-02-11T23:38:08Z

runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.h

@@ -132,6 +117,23 @@ struct PhysPortAndID {
  TUPLE_LIKE_STRUCT_RELATIONAL_OPS(PhysPortAndID)
 };

+struct RouterImpl;
+struct Router {


The struct Router has been moved after the definition of PhysPort to accommodate the newly introduced addFixedPacketConnection function.

jtuyls · 2025-02-12T11:53:41Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

+      slavePorts, packetFlows, priorSlavePorts, priorPacketFlows,
+      deviceModel.getPacketIdMaskWidth());
+
+  // Erase any duplicate packet rules in exising, previously-routed packet


How does this come up? Can we somehow avoid it?

An example is like this:

During the first router run, only the control flow (id=0) is routed:

aie.packet_rules(SOUTH : 1) { aie.rule(31, 0, %[[AMSEL_0]]) {packet_ids = array<i32: 0>} }

In the second router run, when the data flow (id=1) is routed, the packet_rules are updated to:

aie.packet_rules(SOUTH : 1) { aie.rule(30, 0, %[[AMSEL_0]]) {packet_ids = array<i32: 0, 1>} }

Since the control and data packet flows share the same channel (SOUTH:1) and arbiter (%[[AMSEL_0]]), the mask value is updated from 31 to 30.

A probably "safer" option is to replace individual aie.rule entries, rather than erase the aie.packet_rules and regenerate the whole block, but an "issue" is the orders of rules might also get changed due to #1087.

I see, thanks, this makes sense now.

jtuyls · 2025-02-12T12:11:24Z

compiler/plugins/target/AMD-AIE/aie/test/unit_ctrl_routing.mlir

@@ -1,79 +1,402 @@
-// RUN: iree-opt --amdaie-create-pathfinder-flows %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(aie.device(amdaie-create-pathfinder-flows{route-ctrl=true route-data=false},amdaie-create-pathfinder-flows{route-ctrl=false route-data=true}))" --split-input-file %s | FileCheck %s


Maybe rename the file to unit_ctrl_and_data_routing.mlir.

jtuyls · 2025-02-12T12:16:30Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

+                                  Port{packetRulesOp.getSourceBundle(),
+                                       packetRulesOp.getSourceChannel()},
+                                  PhysPort::Direction::SRC};
+        if (slaveGroups.count(sourcePhyPort))


Wait isn't slaveGroups a std::vector<std::vector<PhysPortAndID>>? What am I missing?

After #1087, the data type is now

using SlaveGroupsT = std::map<PhysPort, SmallVector<std::set<uint32_t>>>;

Right, thanks, I still had an older version locally.

jtuyls · 2025-02-12T12:20:29Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

+    // Infer the connections between source and destination ports, by matching
+    // the amsel.
+    for (const auto &[amsel, srcPhysPortAndIDs] : amselToSrcPhysPortAndIDs) {
+      auto &destPhysPorts = amselToDestPhysPorts[amsel];


Suggested change

auto &destPhysPorts = amselToDestPhysPorts[amsel];

SmallVector<PhysPort> &destPhysPorts = amselToDestPhysPorts[amsel];

jtuyls · 2025-02-12T12:20:52Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

+      for (const auto &srcPhysPortAndID : srcPhysPortAndIDs) {
+        for (const auto &destPhysPort : destPhysPorts) {


Use explicit types.

jtuyls · 2025-02-12T12:23:50Z

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

+  // `packet_rules` and `masterset` operations.
+  PacketFlowMapT priorPacketFlows;
+  SmallVector<PhysPortAndID> priorSlavePorts;
+  auto result = getRoutedPacketFlows(device, deviceModel);


Suggested change

auto result = getRoutedPacketFlows(device, deviceModel);

FailureOr<std::tuple<PacketFlowMapT, SmallVector<PhysPortAndID>>> result = getRoutedPacketFlows(device, deviceModel);

jtuyls · 2025-02-12T12:30:34Z

runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.h

+  void initialize(const AMDAIEDeviceModel &targetModel);
+  void addFlow(TileLoc srcCoords, Port srcPort, TileLoc dstCoords, Port dstPort,
+               bool isPacketFlow);
+  bool addFixedCircuitConnection(


We could probably have this operate on PhysPort and avoid passing col/row, similar to addFixedPacketConnection? Not necessarily for this PR.

Yeah, I agree. I will leave it to a different PR.

Another thing, though it might not be a priority now. I think the data structures inside the router could use some cleanup. For example, PathEndPoint is basically the same as PhysPort.

iree-amd-aie/runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.h

Lines 74 to 84 in 8d01ec9

struct PathEndPoint {

TileLoc tileLoc;

Port port;

PathEndPoint() = default;

PathEndPoint(int col, int row, Port port) : PathEndPoint({col, row}, port) {}

PathEndPoint(TileLoc tileLoc, Port port) : tileLoc(tileLoc), port(port) {}

using TupleType = std::tuple<TileLoc, Port>;

PathEndPoint(TupleType t) : PathEndPoint(std::get<0>(t), std::get<1>(t)) {}

operator TupleType() const { return {tileLoc, port}; }

TUPLE_LIKE_STRUCT_RELATIONAL_OPS(PathEndPoint)

};

iree-amd-aie/runtime/src/iree-amd-aie/aie_runtime/iree_aie_router.h

Lines 111 to 123 in 8d01ec9

struct PhysPort {

enum Direction { SRC, DST };

TileLoc tileLoc;

Port port;

Direction direction;

PhysPort(TileLoc t, Port p, Direction direction)

: tileLoc(t), port(p), direction(direction) {}

using TupleType = std::tuple<TileLoc, Port, Direction>;

PhysPort(TupleType t)

: PhysPort(std::get<0>(t), std::get<1>(t), std::get<2>(t)) {}

operator TupleType() const { return {tileLoc, port, direction}; }

TUPLE_LIKE_STRUCT_RELATIONAL_OPS(PhysPort)

};

Yeah, makes sense.

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

jtuyls

Some more nits, but otherwise this looks good to me, thanks!

compiler/plugins/target/AMD-AIE/aie/AMDAIECreatePathFindFlows.cpp

Yu-Zhewen added 8 commits February 6, 2025 00:10

init commit

a677fe5

fix test

282596d

Merge branch 'main' into zhewen_priority_routing

d04e59f

Merge branch 'main' into zhewen_priority_routing

49b8357

fix format

6bdfa15

Merge branch 'main' into zhewen_priority_routing

5f0ff48

test different orders

e823736

Merge branch 'main' into zhewen_priority_routing

6d163c6

Yu-Zhewen changed the title ~~[WIP][Do Not Review] Separate the routing for CTRL and non-CTRL~~ [WIP][Do Not Review] Separate the routing for control flows and data flows Feb 11, 2025

add more comments and cleanup

80f28d1

Yu-Zhewen changed the title ~~[WIP][Do Not Review] Separate the routing for control flows and data flows~~ Separate the routing for control flows and data flows Feb 11, 2025

Yu-Zhewen commented Feb 11, 2025

View reviewed changes

Yu-Zhewen added 2 commits February 12, 2025 00:09

rename pre-routed to prior/previously-routed

7644c29

Merge branch 'main' into zhewen_priority_routing

75dcf00

Yu-Zhewen marked this pull request as ready for review February 12, 2025 00:11

Yu-Zhewen requested review from makslevental, MaheshRavishankar, nirvedhmeshram, yzhang93, Abhishek-Varma and jtuyls as code owners February 12, 2025 00:11

jtuyls reviewed Feb 12, 2025

View reviewed changes

Yu-Zhewen added 2 commits February 12, 2025 16:09

rename test and use explicit types

23b93ab

Merge branch 'main' into zhewen_priority_routing

13eb2f1

Yu-Zhewen force-pushed the zhewen_priority_routing branch from 1fe1ffa to 13eb2f1 Compare February 12, 2025 16:09

jtuyls approved these changes Feb 12, 2025

View reviewed changes

resolve comments

5acf759

Yu-Zhewen enabled auto-merge (squash) February 12, 2025 16:46

Yu-Zhewen merged commit 8505ad7 into main Feb 12, 2025
7 checks passed

Yu-Zhewen deleted the zhewen_priority_routing branch February 12, 2025 17:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Separate the routing for control flows and data flows #1082

Separate the routing for control flows and data flows #1082

Yu-Zhewen commented Feb 6, 2025 •

edited

Loading

Yu-Zhewen Feb 11, 2025

Yu-Zhewen Feb 11, 2025

Yu-Zhewen Feb 11, 2025

Yu-Zhewen Feb 11, 2025 •

edited

Loading

jtuyls Feb 12, 2025

Yu-Zhewen Feb 12, 2025 •

edited

Loading

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

Yu-Zhewen Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls Feb 12, 2025

Yu-Zhewen Feb 12, 2025

jtuyls Feb 12, 2025

jtuyls left a comment

		@@ -1,79 +1,402 @@
		// RUN: iree-opt --amdaie-create-pathfinder-flows %s \| FileCheck %s
		// RUN: iree-opt --pass-pipeline="builtin.module(aie.device(amdaie-create-pathfinder-flows{route-ctrl=true route-data=false},amdaie-create-pathfinder-flows{route-ctrl=false route-data=true}))" --split-input-file %s \| FileCheck %s

	auto &destPhysPorts = amselToDestPhysPorts[amsel];
	SmallVector<PhysPort> &destPhysPorts = amselToDestPhysPorts[amsel];

		for (const auto &srcPhysPortAndID : srcPhysPortAndIDs) {
		for (const auto &destPhysPort : destPhysPorts) {

	auto result = getRoutedPacketFlows(device, deviceModel);
	FailureOr<std::tuple<PacketFlowMapT, SmallVector<PhysPortAndID>>> result = getRoutedPacketFlows(device, deviceModel);

	struct PathEndPoint {
	TileLoc tileLoc;
	Port port;
	PathEndPoint() = default;
	PathEndPoint(int col, int row, Port port) : PathEndPoint({col, row}, port) {}
	PathEndPoint(TileLoc tileLoc, Port port) : tileLoc(tileLoc), port(port) {}
	using TupleType = std::tuple<TileLoc, Port>;
	PathEndPoint(TupleType t) : PathEndPoint(std::get<0>(t), std::get<1>(t)) {}
	operator TupleType() const { return {tileLoc, port}; }
	TUPLE_LIKE_STRUCT_RELATIONAL_OPS(PathEndPoint)
	};

	struct PhysPort {
	enum Direction { SRC, DST };
	TileLoc tileLoc;
	Port port;
	Direction direction;
	PhysPort(TileLoc t, Port p, Direction direction)
	: tileLoc(t), port(p), direction(direction) {}
	using TupleType = std::tuple<TileLoc, Port, Direction>;
	PhysPort(TupleType t)
	: PhysPort(std::get<0>(t), std::get<1>(t), std::get<2>(t)) {}
	operator TupleType() const { return {tileLoc, port, direction}; }
	TUPLE_LIKE_STRUCT_RELATIONAL_OPS(PhysPort)
	};

Separate the routing for control flows and data flows #1082

Separate the routing for control flows and data flows #1082

Conversation

Yu-Zhewen commented Feb 6, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Yu-Zhewen Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Yu-Zhewen Feb 12, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jtuyls left a comment

Choose a reason for hiding this comment

Yu-Zhewen commented Feb 6, 2025 •

edited

Loading

Yu-Zhewen Feb 11, 2025 •

edited

Loading

Yu-Zhewen Feb 12, 2025 •

edited

Loading